/*
 ============================================================================
 Name        : dnsClient.c
 Author      : Brad Lajoie and Adam Blackwell
 Version     :
 Copyright   : Your copyright notice
 Description : Simple DNS Client
 ============================================================================
 */

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <sys/types.h>

#include "dnsStructs.h"
#include "defaults.h"
#include "dnsClientPrototypes.h"

// Checks.
#define pAssert(a,b) {if (!(a)) {printf("Line %d in File %s: %s\n", __LINE__, __FILE__, b); exit(1);}}
#define pError(e,a) {if (e) {perror(a); exit(1);}}

/*
 * dnsclient [-t<timeout>] [-i<max-retries>] [-p<port>] @<server> <name>
 *
 * timeout (Optional) Indicates in milliseconds,
 * how long to wait before regenerating an unanswered query. Default value : 5.
 *
 * max-retries (Optional) Indicates the number of times the resolver will
 * re-generate the query, before it quits with an error,
 * if no response is received from the server. Default value : 3.
 *
 * port (Optional) The UDP port number of the DNS server.  Default value : 53.
 *
 * server (Required) The IP address of the DNS server,in a.b.c.d format.

 * name (Required) The name to query for. e.g. www.google.com
 */

int main(int argc, char *argv[])
{
  /****************
   * Fields:      *
   ****************/
  
  // Query.
  struct Query myQuery;
  initializeQueryStruct(&myQuery);
  
  // Argument Fields.
  int timeout = TIMEOUT;
  int maxRetries = MAX_RETRIES;
  // Retries Field.
  int numRetries = 0;
  int port = PORT_NUMBER;
  char* theServer;
  char* theURL;

  // Socket Fields;
  int socketFileDesc;
  struct sockaddr_in socketIn;
  
  /********************
   * Parse Arguments: *
   ********************/
   
  if(argc < 3)
  printUsageAndExit();
  // Get the optional items
  int i;
  for(i = 1; i < argc; i++) { // skip argc 0, program name
    if(i + 1 == argc)
      continue;

    if(strcmp(argv[i], "-t") == 0) {
      timeout = atoi(argv[i + 1]);
    }
    else if(strcmp(argv[i], "-i") == 0) {
      maxRetries = atoi(argv[i + 1]);
    }
    else if(strcmp(argv[i], "-p") == 0) {
      port = atoi(argv[i + 1]);
    }
  }

  // Get the required options.
  // Get the DNS Server IP.
  theServer = argv[argc - 2];
  char* tok = strtok(theServer, "@");
  theServer = tok;

  // Get the URL Name.
  theURL = argv[argc - 1];

  // Put theURL into qname of Question struct in the Query struct (e.g. www.google.com -> 3www6google3com)
  urlToQname(theURL, &myQuery); 
  // Generate a header ID
  short id = (short)rand() % 65535;
  // Put ID into header
  myQuery.header.id = id;
  
  /********************************
   * Create and Connect Socket:   *
   ********************************/

  // Setup the socket NOTE: for in_addr do inet_aton(theServer)
  socketFileDesc = socket(AF_INET, SOCK_DGRAM, 0); // PF_INET?
  pError((socketFileDesc == -1), "socket");

  bzero((char *) &socketIn, sizeof(socketIn));
  socketIn.sin_family = PF_INET;
  //theServer = "129.10.112.250"; // bogus ip for testing select() and timeout
  inet_aton(theServer, &socketIn.sin_addr);
  socketIn.sin_port = htons(port);

  int rCode = connect(socketFileDesc, (struct sockaddr *)&socketIn, sizeof(socketIn));
  while(rCode != 0) {
    numRetries++;
    printf("Error connecting to server. Retrying %d out of %d...\n", numRetries, maxRetries);
    if(numRetries == maxRetries) {
      printf("ERROR:\tReached the maximum amount of retries, %d. Now exiting.\n", numRetries);
      exit(0);
    }
    rCode = connect(socketFileDesc, (struct sockaddr *)&socketIn, sizeof(socketIn));
  }
  // reset numRetries for later retries on other issues we may encounter
  numRetries = 0;

  /******************
   * Serialization: *
   *****************/
  
  int qnameSize = (sizeof(char) * (strlen(myQuery.question.qname) + 1));
  int questionSize = QUESTION_INFO_SIZE + qnameSize;
  int packetSize = HEADER_INFO_SIZE + questionSize;

  unsigned short* serializedQuery = (unsigned short*) malloc(packetSize);
  zeroMemory(serializedQuery, packetSize);
  unsigned short* mySerializedHeader = (unsigned short*) malloc (HEADER_INFO_SIZE);
  zeroMemory(mySerializedHeader, HEADER_INFO_SIZE);
  unsigned short* mySerializedQuestion = (unsigned short*) malloc (questionSize);
  zeroMemory(mySerializedQuestion, questionSize);

  // Serialize the header and question
  serializeHeader(mySerializedHeader, myQuery.header);
  serializeQuestion(mySerializedQuestion, myQuery.question, questionSize, qnameSize);

  // Concat the serialized header and question to serialized query
  memcpy(serializedQuery, mySerializedHeader, HEADER_INFO_SIZE / 2);
  // Testing serialization and concat of header:
  /*unsigned char *tmp = (unsigned char *) serializedQuery;
  for(i = 0; i < NUM_HEADER_FIELDS * 2; i++) {
    unsigned char field = tmp[i];
    printf("%d : %x\n", field, field);
  }*/
  // offset by header size / 2 because HEADER_INFO_SIZE is in chars, need it in shorts (= 2 * chars)
  memcpy(serializedQuery + (HEADER_INFO_SIZE / 2), mySerializedQuestion, questionSize);
  // Testing serialization of query (and concat of question)
  /*unsigned char *tmp2 = (unsigned char *) serializedQuery;
  for(i = 0; i < packetSize; i++) {
    unsigned char field = tmp2[i];
    printf("%d : %x\n", field, field);
  }*/

  /********************
   * Send Message:    *
   ********************/
   
  int retVal = send(socketFileDesc, (char *)serializedQuery, packetSize, 0);
  while(retVal < 0) {
    numRetries++;
    printf("Error sending packet to server. Retrying %d out of %d...\n", numRetries, maxRetries);
    if(numRetries == maxRetries) {
      printf("ERROR:\tReached the maximum amount of retries, %d. Now exiting.\n", numRetries);
      exit(0);
    }
    retVal = send(socketFileDesc, serializedQuery, packetSize, 0);
  }
  // Reset numRetries for later retries on other issues we may encounter
  numRetries = 0;
  
  /********************
   * Timeout Logic:   *
   ********************/

  // Timeout Struct.
  struct timeval myTimeval;
  myTimeval.tv_sec = timeout;
  myTimeval.tv_usec = 0;

  // File Descriptor Set Up.
  fd_set fds;
  FD_ZERO(&fds);
  FD_SET(socketFileDesc, &fds);
  // Attempt to get a response.
  while(numRetries != maxRetries) {
    // Use selct to manage the reties.
    int selectOutput = select(socketFileDesc + 1, &fds, NULL, NULL, &myTimeval);
    // If there is an error or no response.
    if (selectOutput < 1) {
      // Increment Retries.
      numRetries++;
      // Print Error.
      printf("Error or timeout waiting for response. Retrying %d out of %d...\n",
             numRetries, maxRetries);
      // Try again.
      send(socketFileDesc, serializedQuery, packetSize, 0);
    }
    // Else if a response is recieved.
    else {
      break;
    }
  }

  // If it hit the Max Reties and didn't break
  if(numRetries == maxRetries) {
    // Print an error.
    printf("ERROR:\tReached the maximum amount of retries, %d. Now exiting.\n", numRetries);
    exit(0);
  } 
  // Reset numRetries for later retries on other issues we may encounter.
  numRetries = 0;

  // Check. // TODO: loop retries?
  if (!FD_ISSET(socketFileDesc, &fds)) {
    // Print an Error.
    printf("ERROR:\tWaiting for socket read failed. Now exiting.\n");
    exit(0);
  }

  // TODO: have a retry() procedure instead of all of these while loops for retry. Simply if(error) retry()

  /*********************
   * Receive response: *
   *********************/

  unsigned char* responsePacket = (unsigned char *) malloc(MAX_PACKET_SIZE);
  zeroMemory(responsePacket, MAX_PACKET_SIZE);
  int recvOutput = recv(socketFileDesc, responsePacket, MAX_PACKET_SIZE, 0);
  while(recvOutput < 1) {
    numRetries++;
    // if -1, error.
    // if 0, the socket is blocking and the connection to the remote node failed.
    printf("Error trying to receive packet. Retrying %d out of %d...\n", numRetries, maxRetries);
    if(numRetries == maxRetries) {
      printf("ERROR:\tReached the maximum amount of retries, %d. Now exiting.\n", numRetries);
      exit(0);
    }

    recvOutput = recv(socketFileDesc, responsePacket, MAX_PACKET_SIZE, 0);
  }
  // Testing response packet data
  /*unsigned char *tmp3 = (unsigned char *) responsePacket;
  // packetSize was size of query packet, add on to print out answer (guessing name and rdata length)
  int nameLengthGuess = qnameSize; // NOTE: could be a pointer
  int rdataLengthGuess = 20;
  for(i = 0; i < packetSize + ANSWER_INFO_SIZE + nameLengthGuess + rdataLengthGuess; i++) {
    unsigned char field = tmp3[i];
    printf("%d : %x\n", field, field);
  }*/

  /*********************
   * Deserialization:  *
   *********************/
  
  struct Response myResponse;
  zeroMemory(&myResponse, sizeof(struct Response));
  deserializePacket(&myResponse, responsePacket, questionSize, qnameSize);

  /*********************************
   * Checking validity of response *
   *********************************/
  errorCheckResponse(myResponse, myQuery);

  /*********************
   * Print out to User:*
   *********************/

  // need to loop if multiple answers
  // rdata dependent on type, test type and print out based on type (A record = IP, CNAME = name of alias)
  if(myResponse.answer.type == A_RECORD) {
    printOutARecordResponse(myResponse);
  }
  else if(myResponse.answer.type == CNAME) {
    printOutCNAMEResponse(myResponse);
  }

  return 0;
}

/*********************
 * Helper Functions: *
 *********************/

// Assigns default values to a query struct.  
void initializeQueryStruct(struct Query *myQuery) {
  // Set all bits in myQuery to 0 to begin with
  zeroMemory((void*)myQuery, sizeof(struct Query));

  // Assign initial values to header components (that aren't 0)
  myQuery->header.rd = 1;
  myQuery->header.qdcount = 1;

  // Assign initial values to question components (that aren't 0)
  myQuery->question.qtype = A_RECORD;
  myQuery->question.qclass = A_RECORD;
}

// Zeros out a given pointer to a given length.
void zeroMemory(void* ptr, size_t size)
{
  if (ptr != NULL) {
    memset(ptr, 0, size);
  }
}

// Print Usage and Exit.
void printUsageAndExit() {
  printf("Invalid input. Expected:\n dnsClient [-t <timeout>] [-i <max-retries>] [-p <port>] @<server> <name>\n");
  exit(0);
}

void printOutARecordResponse(struct Response myResponse) {
  printf("IP\t%s\t%d\t%s\n",
      myResponse.answer.rdata,
      myResponse.answer.ttl,
      myResponse.header.aa ? "auth" : "nonauth");
}

void printOutCNAMEResponse(struct Response myResponse) {
  printf("CNAME\t%s\t%d\t%s\n",
      myResponse.answer.rdata,
      myResponse.answer.ttl,
      myResponse.header.aa ? "auth" : "nonauth");
}

// Takes a Serialized query and a URL and formats the URL to be put at the
// front of the query.
void urlToQname(char* theUrl, struct Query* myQuery) {
  char* tok = strtok(theUrl, ".");
  //char buf[1000];
  char* buf = (char*) malloc(1000);
  bzero(buf, 1000);

  int index = 0;
  while (tok != NULL) {
    int count = strlen(tok);

    buf[index] = count;
    index++;

    // Put tok into buf
    int i;
    for(i = 0; i < count; i++) {
      buf[index + i] = (*tok);
      tok++;
    }

    index += i;
    tok = strtok(NULL,".");
  }
  buf[index + 1] = 0;

  myQuery->question.qname = buf;
}

void errorCheckResponse(struct Response myResponse, struct Query myQuery) {
  // Header Checks:
  // verify header ID
  if(myResponse.header.id != myQuery.header.id) {
    printf("ERROR:\tHeader ID in response packet [%d] not the same as query packet [%d]. Exiting\n",
        myResponse.header.id, myQuery.header.id);
    exit(0);
  }
  // if rcode != 0 (switch and print out messages explaining rcodes 1-5 then exit, & treat 3 special)
  if(myResponse.header.rcode) {
    switch(myResponse.header.rcode) {
      // format error
      case 1:
        printf("ERROR:\tFormat error. Name server unable to interpret the query, exiting.\n");
        break;
      // server failure
      case 2:
        printf("ERROR:\tServer failure. Name server was unable to process the query due to a problem with the "
            "name server, exiting.\n");
        break;
      // name error
      case 3:
        printf("NOTFOUND:\tDomain name referenced in the query [%s] does not exist, exiting.\n",
            myQuery.question.qname);
        break;
      // query not supported by server
      case 4:
        printf("ERROR:\tThe name server does not support this kind of query, exiting\n");
        break;
      // server refused query
      case 5:
        printf("ERROR:\tThe name server refuses to perform the specified operation for policy reasons, exiting\n");
        break;
      // unexpect rcode
      default:
        printf("ERROR:\tHow in the hell did you get here, rcode [%d] does not exist! Exiting\n",
            myResponse.header.rcode);
        break;
    }
    exit(0);
  }
  // expect qr = 1
  if(!myResponse.header.qr) {
    printf("ERROR:\tPacket received is not a response, qr = %d. Exiting\n", myResponse.header.qr);
    exit(0);
  }
  // test tc = 1 then exit
  if(myResponse.header.tc) {
    printf("ERROR:\tResponse is truncated, exiting.\n");
    exit(0);
  }
  // if ra = 0 exit
  if(!myResponse.header.ra) {
    printf("ERROR:\tServer does not support recursion, exiting.\n");
    exit(0);
  }
  // ignore nscount
  // ignore arcount

  // Question Checks:
  // qtype = 0x0001 (A Record), else exit
  if(myResponse.question.qtype != A_RECORD) {
    printf("ERROR:\tResponse question's qtype is not an A Record, exiting.\n");
    exit(0);
  }
  // qclass = 0x0001, else exit
  if(myResponse.question.qclass != A_RECORD) {
    printf("ERROR:\tResponse question's qclass is not an A Record, exiting.\n");
    exit(0);
  }

  // Answer Checks:
  // type = 0x0001 (A Record); = 0x0005 (CNAME), accept both, others exit
  if(!(myResponse.answer.type == A_RECORD || myResponse.answer.type == CNAME)) {
    printf("ERROR:\tResponse answer type [%d] not supported. Only support A Records [%d] and CNAMEs [%d]. Exiting\n",
        myResponse.answer.type, A_RECORD, CNAME);
    exit(0);
  }
  // class = 0x0001, else exit
  if(myResponse.answer.theClass != A_RECORD) {
    printf("ERROR:\tResponse class [%d] not valid. Expecting [%d]. Exiting\n",
        myResponse.answer.theClass, A_RECORD);
    exit(0);
  }
}

/***********************
 * Serialization Work: *
 ***********************/

// Takes all values in the Header Struct and creates one Unsigned Short Array
void serializeHeader(unsigned short* head, struct Header header) {
  head[0] = htons(header.id);
  head[1] |= (((unsigned short) header.qr) << 15);
  head[1] |= (((unsigned short) header.opCode) << 11);
  head[1] |= (((unsigned short) header.aa) << 10);
  head[1] |= (((unsigned short) header.tc) << 9);
  head[1] |= (((unsigned short) header.rd) << 8);
  head[1] |= (((unsigned short) header.ra) << 7);
  head[1] |= (((unsigned short) header.z) << 4);
  head[1] |= ((unsigned short) header.rcode);
  head[1] = htons(head[1]);
  head[2] = htons(header.qdcount);
  head[3] = htons(header.ancount);
  head[4] = htons(header.nscount);
  head[5] = htons(header.arcount);

  // Test serialization of header
  /*int i;
  unsigned char* tmp = (unsigned char *) head;
  for(i = 0; i < HEADER_INFO_SIZE; i++) {
    unsigned char field = tmp[i];
    printf("%d : %x\n", field, field);
  }*/
}

// Takes all values in the Question Struct and creates one Unsigned Short Array.
void serializeQuestion(unsigned short* quest, struct Question question, int questionSize, int qnameSize) {
  unsigned short qtypeNetworkOrder = htons(question.qtype);
  unsigned short qclassNetworkOrder = htons(question.qclass);

  unsigned char* tmp = (unsigned char*) quest;
  memcpy(tmp, question.qname, qnameSize);
  memcpy(tmp + qnameSize, &qtypeNetworkOrder, 2);
  memcpy(tmp + qnameSize + 2, &qclassNetworkOrder, 2);

  memcpy(quest, tmp, questionSize);

  // Test serialization of question
  /*int i;
  unsigned char* tmp2 = (unsigned char *) quest;
  for(i = 0; i < questionSize; i++) {
    unsigned char field = tmp2[i];
    printf("%d : %x\n", field, field);
  }*/
}


/*************************
 * Deserialization Work: *
 *************************/

// Casting to Shorts and back caused issues, so start with chars, and cats to
// shorts for the Header.
// Takes all values in the Packet Unsigned Short Array and builds a Query Struct.
void deserializePacket(struct Response* response, unsigned char* pack, int questionSize, int qnameSize) {
  // Sub Arrays:
  unsigned short* head = (unsigned short*)pack;
  unsigned char* quest = pack + HEADER_INFO_SIZE;
  // TODO: questionSize may change because of pointer in qname? (compression)
  unsigned char* answ = pack + HEADER_INFO_SIZE + questionSize;

  // Split up the Work.
  deserializeHeader(&response->header, head);
  deserializeQuestion(&response->question, quest, questionSize, qnameSize);
  deserializeAnswer(&response->answer, answ, pack, qnameSize);
}

// Takes all values in the Packet Unsigned Short Array and builds a Query Struct.
void deserializeHeader(struct Header* header, unsigned short* head) {
  // Testing header portion of received packet
  /*int i;
  unsigned char *tmp = (unsigned char *) head;
  for(i = 0; i < HEADER_INFO_SIZE; i++) {
    unsigned char field = tmp[i];
    printf("%d : %x\n", field, field);
  }*/

  header->id = ntohs(head[0]);
  head[1] = ntohs(head[1]);
  header->qr = (unsigned char)(head[1] >> 15);
  header->opCode = (unsigned char)(head[1] >> 11);
  header->aa = (unsigned char)(head[1] >> 10);
  header->tc = (unsigned char)(head[1] >> 9);
  header->rd = (unsigned char)(head[1] >> 8);
  header->ra = (unsigned char)(head[1] >> 7);
  header->z = (unsigned char)(head[1] >> 4);
  header->rcode = (unsigned char)(head[1]);

  header->qdcount = ntohs(head[2]);
  header->ancount = ntohs(head[3]);
  header->nscount = ntohs(head[4]);
  header->arcount = ntohs(head[5]);
}

// Takes all values in the Question Unsigned Short Array and builds a Question Struct.
void deserializeQuestion(struct Question* question, unsigned char* quest, int questionSize, int qnameSize) {
  // TODO: could quest qname be a pointer? handle this (test for 11 at beginning)
  memcpy(&question->qname, quest, qnameSize);

  unsigned short* qtype_p = (unsigned short*)(quest + qnameSize);
  unsigned short qtype = ntohs(*qtype_p);
  memcpy(&question->qtype, &qtype, 2);

  unsigned short* qclass_p = (unsigned short*)(quest + qnameSize + 2);
  unsigned short qclass = ntohs(*qclass_p);
  memcpy(&question->qclass, &qclass, 2);
}

<<<<<<< .mine
 

=======
>>>>>>> .r93
// Takes all values in the Answer Unsigned Short Array and builds an Answer Struct.
void deserializeAnswer(struct Answer* answer, unsigned char* answ, unsigned char* pack, int qnameSize) {
  unsigned short nameSize;
  unsigned char firstByte = answ[0];
  if((unsigned short)firstByte >= 192) {
    nameSize = 2; // 2 bytes for a pointer
  }
  else {
    nameSize = qnameSize; // name size of name field in bytes is same as qname size
  }
  answer->name = deserializeName(answ, &firstByte, pack);

  unsigned short* type_p = (unsigned short*)(answ + nameSize);
  unsigned short type = ntohs(*type_p);
  memcpy(&answer->type, &type, 2);

  unsigned short* theClass_p = (unsigned short*)(answ + nameSize + 2);
  unsigned short theClass = ntohs(*theClass_p);
  memcpy(&answer->theClass, &theClass, 2);

  unsigned short* ttl_p = (unsigned short*)(answ + nameSize + 6);
  unsigned short ttl = ntohs(*ttl_p);
  memcpy(&answer->ttl, &ttl, 2);

  unsigned short* rdlength_p = (unsigned short*)(answ + nameSize + 8);
  unsigned short rdlength = ntohs(*rdlength_p);
  memcpy(&answer->rdlength, &rdlength, 2);

  // TODO: use rdlength with rdata? (could have more than one name/ip)
  firstByte = answ[nameSize + 10];
  if(type == A_RECORD) {
    answer->rdata = (char*)malloc((IP_SIZE * 2) - 1);
    int i;
    for(i = 0; i < IP_SIZE; i++) {
      unsigned short octet = (unsigned short)answ[nameSize + 10 + i];
      char buf[3];
      sprintf(buf, "%d", octet);
      // append field to rdata
      strcat(answer->rdata, buf);
      if(i + 1 != IP_SIZE) {
        strcat(answer->rdata, ".");
      }
    }
  }
  else if(type == CNAME) {
    answer->rdata = deserializeName(answ + nameSize + 10, &firstByte, pack);
  }
  else {
    printf("ERROR:\tResponse answer type [%d] not supported. Only support A Records [%d] and CNAMEs [%d]. Exiting\n",
        type, A_RECORD, CNAME);
    exit(0);
  }
}
<<<<<<< .mine
=======

char* deserializeName(unsigned char* readStartPosition, unsigned char* firstByte, unsigned char* pack) {
  // read first byte
  // base case: if firstByte is 0, done, return firstByte
  if((unsigned short)*firstByte == 0) {
    return '\0';
  }
  // if firstByte < 192, this part of the name, read this number of bytes, put a "."
  // then recurse (0 will be added in base case);
  else if((unsigned short)*firstByte < 192) {
    unsigned short numBytesToRead = (unsigned short)*firstByte;
    char* sequence = (char*)malloc(numBytesToRead + 1);
    zeroMemory(sequence, numBytesToRead + 2);
    memcpy(sequence, readStartPosition + 1, numBytesToRead);

    readStartPosition += numBytesToRead + 1;
    char* theRest = deserializeName(readStartPosition, &readStartPosition[0], pack);
    if(theRest != 0) {
      strcat(sequence, ".");
      strcat(sequence, theRest);
    }
    return sequence;
  }
  // if firstByte >= 192, pointer. Get the offset (after first 2 bits) by setting the first two bits of
  // firstByte to 0 (& with 0x3F), recurse starting at this offset into the packet
  else {
    unsigned short offset = (*firstByte & 0x3F) + readStartPosition[1];
    readStartPosition = pack + offset;
    return deserializeName(readStartPosition, &readStartPosition[0], pack);
  }
}
